package main

import (
	"fmt"
	"sync"
	"time"
)

// MVCCStore 代表支持MVCC的存储结构
type MVCCStore struct {
	// 存储不同版本的数据，外层map的key是数据的键，内层map的key是版本号，value是数据值
	data map[string]map[int64]string
	// 当前最新版本号
	currentVersion int64
	// 为每个key维护的最新版本号
	latestVersions map[string]int64
	// 读写锁，用于保护并发访问
	mutex sync.RWMutex
}

// NewMVCCStore 创建新的MVCC存储
func NewMVCCStore() *MVCCStore {
	return &MVCCStore{
		data:           make(map[string]map[int64]string),
		currentVersion: 0,
		latestVersions: make(map[string]int64),
	}
}

// Put 设置键值对，创建新版本
func (mvcc *MVCCStore) Put(key, value string) {
	mvcc.mutex.Lock()
	defer mvcc.mutex.Unlock()

	// 递增版本号
	mvcc.currentVersion++
	version := mvcc.currentVersion

	// 如果key不存在，则初始化版本映射
	if _, exists := mvcc.data[key]; !exists {
		mvcc.data[key] = make(map[int64]string)
	}

	// 存储当前版本的值
	mvcc.data[key][version] = value

	// 更新最新版本号
	mvcc.latestVersions[key] = version
}

// Get 获取键的最新值
func (mvcc *MVCCStore) Get(key string) (string, bool) {
	mvcc.mutex.RLock()
	defer mvcc.mutex.RUnlock()

	// 获取key的最新版本号
	version, exists := mvcc.latestVersions[key]
	if !exists {
		return "", false
	}

	// 返回最新版本的值
	value, exists := mvcc.data[key][version]
	return value, exists
}

// GetAtVersion 获取指定版本的值
func (mvcc *MVCCStore) GetAtVersion(key string, version int64) (string, bool) {
	mvcc.mutex.RLock()
	defer mvcc.mutex.RUnlock()

	// 检查key是否存在
	versions, exists := mvcc.data[key]
	if !exists {
		return "", false
	}

	// 获取指定版本的值，如果不存在则尝试获取小于该版本的最大版本的值
	value, exists := versions[version]
	if exists {
		return value, true
	}

	// 寻找小于指定版本的最大版本
	var maxVersion int64 = 0
	for v := range versions {
		if v < version && v > maxVersion {
			maxVersion = v
		}
	}

	if maxVersion > 0 {
		return versions[maxVersion], true
	}

	return "", false
}

// GetCurrentVersion 获取当前版本号
func (mvcc *MVCCStore) GetCurrentVersion() int64 {
	mvcc.mutex.RLock()
	defer mvcc.mutex.RUnlock()
	return mvcc.currentVersion
}

// Snapshot 创建指定版本的快照
type Snapshot struct {
	store   *MVCCStore
	version int64
}

// NewSnapshot 创建当前版本的快照
func (mvcc *MVCCStore) NewSnapshot() *Snapshot {
	return &Snapshot{
		store:   mvcc,
		version: mvcc.GetCurrentVersion(),
	}
}

// Get 从快照中获取值
func (s *Snapshot) Get(key string) (string, bool) {
	return s.store.GetAtVersion(key, s.version)
}

// 演示MVCC的使用
func main() {
	// 创建MVCC存储
	mvcc := NewMVCCStore()

	// 模拟并发写入
	var wg sync.WaitGroup
	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			key := fmt.Sprintf("key%d", id)
			mvcc.Put(key, fmt.Sprintf("value%d-1", id))
			time.Sleep(10 * time.Millisecond)
			mvcc.Put(key, fmt.Sprintf("value%d-2", id))
		}(i)
	}

	// 等待所有写入完成
	wg.Wait()

	// 创建快照
	snapshot := mvcc.NewSnapshot()
	snapshotVersion := snapshot.version
	fmt.Printf("创建了版本 %d 的快照\n", snapshotVersion)

	// 继续更新数据
	for i := 0; i < 3; i++ {
		mvcc.Put(fmt.Sprintf("key%d", i), fmt.Sprintf("value%d-3", i))
	}

	// 模拟并发读取 - 一些读取最新版本，一些读取快照版本
	wg = sync.WaitGroup{}
	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			key := fmt.Sprintf("key%d", id)
			
			// 读取最新版本
			latestValue, exists := mvcc.Get(key)
			if exists {
				fmt.Printf("最新版本: %s = %s\n", key, latestValue)
			}
			
			// 读取快照版本
			snapshotValue, exists := snapshot.Get(key)
			if exists {
				fmt.Printf("快照版本(%d): %s = %s\n", snapshotVersion, key, snapshotValue)
			}
		}(i)
	}

	wg.Wait()

	// 输出版本信息
	fmt.Printf("当前版本: %d\n", mvcc.GetCurrentVersion())
	fmt.Println("各key的最新版本:")
	mvcc.mutex.RLock()
	for k, v := range mvcc.latestVersions {
		fmt.Printf("%s: 版本 %d\n", k, v)
	}
	mvcc.mutex.RUnlock()
}